package com.springboot.kafka;

import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;

import java.util.List;
import java.util.Map;
import java.util.Random;

public class HotDataPartitioner implements Partitioner {

    private Random random;

    @Override
    public void configure(Map<String, ?> configs) {
        random = new Random();
    }

    @Override
    public int partition(String topic, Object keyObj, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
        String key = (String) keyObj;
        List<PartitionInfo> partitionInfoList = cluster.availablePartitionsForTopic(topic);
        int partitionCount = partitionInfoList.size();
        int hotDataPartition = partitionCount - 1;
        return !key.contains("hot_data") ? random.nextInt(partitionCount - 1) : hotDataPartition;
    }

    @Override
    public void close() {

    }

    @Override
    public void onNewBatch(String topic, Cluster cluster, int prevPartition) {

    }

}

// props.put("partitioner.class","com.zhss.HotDataPartitioner");
//
//         测试发送
//
//         bin/kafka-run-class.sh kafka.tools.GetOffsetShell--broker-list localhost:9092--topic test-topic